#include <iostream>
#include <stdio.h>
#include "../time_cost.hpp"
#include "math.h"
#include <stdint.h>
// #include <stdlib.h>
#include <cmath>

#include <chrono>
#include <functional>
#include <fstream>
using namespace std;
int main(int argc, char** argv){
    int N = 10;
    int _use_scale = 1;
    int warm_up_times = 1;
    int test_times = 10;

    if (argc>1) N = atoi(argv[1]);
    if (argc>2) _use_scale = atoi(argv[2]);
    if (argc>3) warm_up_times = atoi(argv[3]);
    if (argc>4) test_times = atoi(argv[4]);

    int* dst;
    float* src;
    int len=1<<N;
    int PACKING_SIZE = sizeof(float)*8;
    int padding_len = (PACKING_SIZE - (len % PACKING_SIZE)) % PACKING_SIZE;
    int chunk_len = (len + padding_len) / PACKING_SIZE;
    src = (float*) calloc(len, sizeof(float));
    dst = (int*) calloc(chunk_len+1, sizeof(float));

    srand(time(0));

    
    auto init_data = [&](){
        for (int i = 0 ; i < len; i++ ){
            src[i] = (float(rand())/RAND_MAX)*2-1;
        }
    };



    printf( "len=%d\n"
            "PACKING_SIZE=%d\n"
            "padding_len=%d\n"
            "chunk_len=%d\n"
            , len, PACKING_SIZE, padding_len, chunk_len);

    auto compress = [&](){
        float scale = 1.0f;
        if (_use_scale) {
            double sum = 0.0f;
            for (int i = 0; i < len; i++){
                sum = sum + abs(src[i]);
            }
            scale = sum / len;
        }

        #pragma omp parallel for simd
        for (int i = 0; i < chunk_len; i++){
            int x = src[i*PACKING_SIZE] < 0;
            for (int j = 1; j < PACKING_SIZE; j++){
                x <<= 1;
                x |= src[i*PACKING_SIZE + j] < 0;
            }
            dst[i] = x;
        }
        float* p_scale = reinterpret_cast<float*>(&dst[chunk_len]);
        *p_scale = scale; 
    };
    auto decompress = [&](){
        auto* pf = reinterpret_cast<float*> (dst+chunk_len);
        float scale = *pf;
        #pragma omp parallel for simd
        for (int i = chunk_len-1; i>=0; --i){
            int x = dst[i];
            for (int j = PACKING_SIZE - 1; j >= 0; j--){
                int sign = 1 - ((x&0x01) << 1);
                src[i*PACKING_SIZE+j] = sign * scale;
                x>>=1;
            }
        }
    };


    int total_compression_time = 0;
    int total_decompression_time = 0;
    auto run = [&](){
        init_data();
        auto start = chrono::system_clock::now();
        compress();
        auto end = chrono::system_clock::now();
        auto duration = chrono::duration_cast<chrono::microseconds>(end-start);
        total_compression_time += duration.count();
        
        start = chrono::system_clock::now();
        decompress();
        end = chrono::system_clock::now();
        duration = chrono::duration_cast<chrono::milliseconds>(end-start);
        total_decompression_time += duration.count();
    };

    for (int i = 0; i < warm_up_times; i++){
        run();
    }
    total_compression_time = total_decompression_time = 0;
    for (int i = 0 ; i < test_times; i++){
        run();
    }
    double average_compression_cost = ((double)total_compression_time)/test_times;
    double average_decompression_cost = ((double)total_decompression_time)/test_times;
    printf("average compression cost %lf us\n", average_compression_cost);
    printf("average decompression cost %lf us\n", average_decompression_cost);
    ofstream out("ret.txt");
    out << average_compression_cost << endl;
    out << average_decompression_cost << endl;
    out.close();
    
    return 0;

}